import pandas as pd
import matplotlib.pyplot as plt
from xgboost import XGBRegressor
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import GridSearchCV
from sklearn.metrics import mean_squared_error, mean_absolute_error
import joblib
from utils.data_util import preprocessing
from utils.log import Logger

plt.rcParams['font.family'] = 'SimHei'
plt.rcParams['font.size'] = 15

# 设置显示的最大行数和列数为None
pd.set_option('display.max_rows', None)
pd.set_option('display.max_columns', None)

# 创建日志对象
logging_obj = Logger(root_path="../",log_name="train")
logger = logging_obj.get_logger()

def init():
    df = preprocessing("../data/train.csv")
    return df

def eda(df):
    logger.info("开始进行数据探索")

    # 1- 创建画布
    fig = plt.figure(figsize=(20,40),dpi=150)

    # 2- 数据探索
    # 2.1- 绘制用电量的直方图
    """
        add_subplot(4,1,1)
            4：画布分为4行
            1：画布分为1列
            1：第一个图形展示的索引位置。注意：索引从1开始
    """
    graph_1 = fig.add_subplot(4,1,1)
    graph_1.hist(df["power_load"],bins=100,rwidth=0.85)
    graph_1.set_title("绘制用电量的直方图")

    # 2.2- 12个月份和负荷存在关系
    """
        month  power_load
        10        1
        10        3
        11        2
        11        4
        11        1
        ....
        
        select month,avg(power_load) from tb group by month
    """
    # 提取月份
    df["month"] = pd.to_datetime(df["time"]).dt.month
    # print(df.head())

    # 按照月份分组，对power_load进行计算均值的聚合操作
    """
        需要设置as_index为False，因为后续plot的时候是将month作为字段使用。否则会报如下的错
        ValueError: 'power_load' is neither a data key nor a valid format string (two marker symbols)
    """
    month_df = df.groupby(["month"],as_index=False)["power_load"].mean()
    # print(month_df)

    # 绘制图形
    graph_2 = fig.add_subplot(4,1,2)
    graph_2.plot("month","power_load",data=month_df)
    graph_2.set_title("12个月份和负荷的关系")

    # 2.3- 24小时和负荷存在关系
    # 提取小时
    df["hour"] = pd.to_datetime(df["time"]).dt.hour
    # print(df.head())

    # 按照小时分组，对power_load进行计算均值的聚合操作
    hour_df = df.groupby(["hour"], as_index=False)["power_load"].mean()

    # 绘制图形
    graph_3 = fig.add_subplot(4, 1, 3)
    graph_3.plot("hour", "power_load", data=hour_df)
    graph_3.set_title("24小时和负荷的关系")

    # 2.4- 节假日和工作日与负荷的关系
    # 特征提取，根据日期时间判断是否是节假日：time先转成日期时间对象 -> 获取是周几 -> 判断是否是节假日
    def is_holiday(val):
        # time先转成日期时间对象 -> 获取是周几
        weekday = pd.to_datetime(val).weekday()

        # 判断是否是节假日
        # 返回值含义：1表示节假日；0表示工作日
        if weekday>=5:
            return 1
        else:
            return 0

    df["is_holiday"] = df["time"].apply(is_holiday)
    # print(df.head(200))

    # 分别计算节假日（周六、周日）、工作日（周一至周五）的用电负荷
    workday_df = df[df["is_holiday"]==0]["power_load"].mean()
    weekday_df = df[df["is_holiday"]==1]["power_load"].mean()

    # 绘制图形
    graph_4 = fig.add_subplot(4,1,4)
    graph_4.bar(x=["工作日","节假日"],height=[workday_df,weekday_df])
    graph_4.set_title("节假日和工作日与负荷的关系")

    fig.show()

    logger.info("完成数据探索")

def feature_extract(df):
    logger.info("开始进行特征提取")

    # 1- 基于time特征得到月份和小时
    df["month"] = pd.to_datetime(df["time"]).dt.month
    df["hour"] = pd.to_datetime(df["time"]).dt.hour

    # 2- 获取昨天同一时刻的用电负荷
    # 2.1- 将数据的格式转成字典。time作为key，power_load作为value
    time_load_dict = df.set_index("time")["power_load"].to_dict()
    # print(type(time_load_dict))
    # print(time_load_dict)

    def get_yesterday_load(row,interval):
        # 2.2- 基于当前日期时间获得昨天同一时刻的日期时间
        # 注意：current_time、yesterday_time不是字符串类型，而是DateTime日期时间对象类型
        current_time = pd.to_datetime(row["time"])
        yesterday_time = current_time - pd.to_timedelta(interval)

        # 2.3- 通过时刻获得负荷
        yesterday_time_str = yesterday_time.strftime("%Y-%m-%d %H:%M:%S")
        return time_load_dict.get(yesterday_time_str)

    # axis=1表示对DataFrame中的数据按行进行逐行处理
    df["yesterday"] = df.apply(get_yesterday_load,axis=1,args=("1d",))

    # 3- 获取当前时刻前1个小时、2个小时、3个小时的用电负荷
    # 实现方式一
    # df["pre_1_hour"] = df.apply(get_yesterday_load,axis=1,args=("1h",))
    # df["pre_2_hour"] = df.apply(get_yesterday_load,axis=1,args=("2h",))
    # df["pre_3_hour"] = df.apply(get_yesterday_load,axis=1,args=("3h",))

    # 实现方式二
    hours = 3
    shift_list = [df["power_load"].shift(i) for i in range(1,hours+1)]
    shift_data = pd.concat(shift_list,axis=1)
    shift_data.columns = [f"pre_{i}_hour" for i in range(1,hours+1)]
    # print(shift_list)
    # print(shift_data)
    df = pd.concat([df,shift_data],axis=1)

    # print(df.head(30))

    logger.info("完成特征提取")

    return df

def train(df):
    logger.info("开始进行模型训练")

    # 1- 拆分得到特征数据和目标值
    """
        特征中需要将time这个字段删除，因为它是非数值类型的字段。如果不删除，会报如下的错误
        ValueError: could not convert string to float: '2014-08-15 13:00:00'
    """
    all_columns = list(df.columns)
    all_columns.remove("power_load")
    all_columns.remove("time")
    x = df[all_columns] # 这里实际是需要写两层中括号，但是因为all_columns本身是一个list列表，因此只需要写一层
    y = df["power_load"]

    # print(all_columns)
    # print(x.head())

    # 2- 划分得到训练集和测试集
    x_train,x_test,y_train,y_test = train_test_split(x,y,test_size=0.2,random_state=107)

    # 3- 特征预处理：标准化处理
    transformer = StandardScaler()
    x_train = transformer.fit_transform(x_train)
    x_test = transformer.transform(x_test)

    # 4- 模型训练
    model = XGBRegressor(n_estimators=100,max_depth=5,random_state=107)
    # model = XGBRegressor()
    # param_dict = {
    #     "n_estimators":[i*100 for i in range(1,4)],
    #     "max_depth":[5*i for i in range(1,4)],
    #     "random_state":[107]
    # }
    # model = GridSearchCV(estimator=model,param_grid=param_dict,cv=2)
    model.fit(x_train,y_train)
    # print("最优超参数组合",model.best_params_)
    # logger.info(f"最优超参数组合{model.best_params_}")


    # 5- 评估
    y_predict = model.predict(x_test)
    # 注意：提交到代码仓库中的代码或者部署到生产环境的代码，不允许出现print打印语句
    # print("模型训练_MSE",mean_squared_error(y_test, y_predict))
    # print("模型训练_MAE",mean_absolute_error(y_test, y_predict))

    logger.info(f"模型训练_MSE{mean_squared_error(y_test, y_predict)}")
    logger.info(f"模型训练_MAE{mean_absolute_error(y_test, y_predict)}")

    # 6- 保存训练好的模型
    joblib.dump(model,"../model/power_load.pkl")
    joblib.dump(transformer,"../model/transformer.pkl")

if __name__ == '__main__':
    # 1- 初始化：数据基本处理
    init_df = init()

    # 2- 数据探索
    # eda(init_df)

    # 3- 特征工程【重点、难点】
    df = feature_extract(init_df)

    # 4- 模型训练
    train(df)
